package org.etlt.load;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.etlt.EtltRuntimeException;
import org.etlt.JavaStyleIdentifier;
import org.etlt.extract.DatabaseUtil;
import org.etlt.extract.Extractor;
import org.etlt.job.JobContext;
import org.etlt.job.PkRange;

import javax.sql.DataSource;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;

public class DatabaseLoader extends Loader {

    protected Log log = LogFactory.getLog(getClass());

    private DataSource dataSource;

    private CountDownLatch countDownLatch;


    public DatabaseLoader(DatabaseLoaderSetting setting) {
        super(setting);
        setName(setting.getName());
    }

    protected void init(JobContext context) {
        DatabaseLoaderSetting setting = getSetting();
        context.createExtractor(setting.getExtractor()).init(context);
        this.dataSource = context.getResource(setting.getDatasource());
        this.countDownLatch = new CountDownLatch(setting.getParallel());
        resolveColumns(context);
    }

    @Override
    public void preLoad(JobContext context) {
        init(context);
        DatabaseLoaderSetting setting = getSetting();
        setting.getPreDmls().forEach(preDml -> {
            try (PreparedStatement statement = this.getConnection().prepareStatement(preDml)) {
                statement.execute();
            } catch (SQLException e) {
                throw new EtltRuntimeException("preLoad error:" + getName(), e);
            }
        });
    }

    public void load(JobContext context) {
        DatabaseLoaderSetting setting = getSetting();
        ThreadGroup threadGroup = new ThreadGroup(getName());
        JobContext newContext = context.copy();

        String ds = setting.getExtractor();
        Extractor extractor = context.createExtractor(ds);
        extractor.init(context);
        if (extractor.getSetting().isParallel()) {// extract parallelled
            PkRange pkRange = extractor.getPkFullRange();
            List<PkRange> pkRanges = pkRange.split(setting.getParallel());
            for (int i = 0; i < Math.min(setting.getParallel(), pkRanges.size()); i++) {
                Extractor innerExtractor = context.createExtractor(ds);
                Thread worker = new Thread(threadGroup,
                        () -> load0(newContext, innerExtractor, getConnection()), getName() + "-worker-" + i);
                worker.start();
            }
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                throw new EtltRuntimeException(e);
            }
        } else {// extract serially
            for (int i = 0; i < setting.getParallel(); i++) {
                Thread worker = new Thread(threadGroup,
                        () -> load0(newContext, extractor, getConnection()), getName() + "-worker");
                worker.start();
            }
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                throw new EtltRuntimeException(e);
            }
        }
    }

    /**
     * support when attribute in dml node
     *
     * @param context    JobContext object
     * @param connection java.sql.Connection object
     */
    private void load0(JobContext context, Extractor extractor, Connection connection) {
        boolean autoCommit = false;
        try {
            DatabaseLoaderSetting setting = getSetting();
            String ds = setting.getExtractor();
            List<Dml> dmls = new ArrayList<>();
            for (DmlSetting dmlSetting : setting.getDmls()) {
                dmls.add(new Dml(dmlSetting));
            }
            autoCommit = connection.getAutoCommit();
            if (autoCommit) connection.setAutoCommit(false);
            for (extractor.extract(context); context.exists(ds); extractor.extract(context)) {
                for (Dml dml : dmls) {
                    for (int i = 0; i < dml.dmlSetting.getRepeat(); i++)
                        dml.execute(context, connection);
                }
            }
            for (Dml dml : dmls) {
                dml.executeTail(connection);
            }

        } catch (Exception e) {
            try {
                if (!connection.getAutoCommit())
                    connection.rollback();
            } catch (SQLException e1) {
                e1.printStackTrace();
            }
            log.warn("executing loader error: " + getName(), e);
            context.fail(this);
        } finally {
            this.countDownLatch.countDown();
            try {
                connection.setAutoCommit(autoCommit);
            } catch (SQLException e) {
                //do nothing
            }
            close(connection);
            if (extractor != null) {
                try {
                    extractor.close();
                } catch (Exception e) {
                    throw new EtltRuntimeException(e);
                }
            }
        }
    }


    protected boolean when(String condition, JobContext context) {
        if (StringUtils.isBlank(condition))
            return true;
        return Boolean.parseBoolean(String.valueOf(context.getExpressionEvaluator().evaluate(condition, context)));
    }

    public void doFinish(JobContext context) {
        context.finish(this);
    }

    protected Connection getConnection() {
        try {
            return this.dataSource.getConnection();
        } catch (SQLException e) {
            throw new EtltRuntimeException(e);
        }
    }

    @Override
    public void close() throws Exception {

    }

    private class Dml {
        private final DmlSetting dmlSetting;

        private String generatedSql;

        private String valueSegment;

        private String originalSql;

        private PreparedStatement preparedStatement;

        private Statement statement;

        private int batched = 0;

        private final Map<String, Object> entityContainer = new HashMap<>();

        public Dml(DmlSetting dmlSetting) {
            this.dmlSetting = dmlSetting;
            this.dmlSetting.parse();
        }

        /**
         * @param context
         * @param connection
         * @throws SQLException
         */
        public void execute(JobContext context, Connection connection) throws SQLException {
            populateEntity(context);
            if (preparedStatement == null) {
                Extractor extractor = context.getExtractor(DatabaseLoader.this.getSetting().getExtractor());
                this.originalSql = dmlSetting.getStandardSql(extractor.getColumns().size());
                this.preparedStatement = connection.prepareStatement(originalSql);
            }
            if (when(dmlSetting.getWhen(), context)) {
                if (!this.dmlSetting.isMultiValue()) {
                    if(log.isDebugEnabled()){
                        log.debug("{" + this.dmlSetting.getId() + "} " + originalSql);
                    }
                    execute_preparedSql(context, connection);
                } else {
                    execute_generatedSql(context, connection);
                }
            }
        }

        private void populateEntity(JobContext context) {
            List<DmlEntitySetting> dmlEntitySettings = this.dmlSetting.getDmlEntitySettings();
            dmlEntitySettings.forEach(dmlEntitySetting -> {
                Object value =
                        context.getExpressionEvaluator().evaluate(dmlEntitySetting.getExpression(), context);
                entityContainer.put(dmlEntitySetting.getName(), value);
            });
        }

        private void createGeneratedSql() {
            if (StringUtils.isEmpty(generatedSql)) {
                String formattedSql = originalSql.toUpperCase();
                String values = " VALUES ";
                int index = formattedSql.indexOf(values);
                if (index == -1) {
                    values = " VALUES(";
                    index = formattedSql.indexOf(values);
                    if(index == -1)
                        throw new IllegalArgumentException(
                                "this sql doesn't support multi value statement, ' VALUES ' or ' VALUES(' should be included: " + this.originalSql);
                }
                generatedSql = formattedSql.substring(0, index).concat(" VALUES");
                if (StringUtils.isEmpty(valueSegment)) {
                    valueSegment = formattedSql.substring(generatedSql.length());
                }
            }
        }

        /**
         * handle all entity
         * @param context
         * @param expression
         * @return
         */
        private String preEvaluate(JobContext context, String expression){
            StringBuilder stringBuilder = new StringBuilder();
            int start = 0;
            for(int index = expression.indexOf(DmlSetting.ENTITY_IDENTIFIER, start); index > -1;
                index = expression.indexOf(DmlSetting.ENTITY_IDENTIFIER, start)){
                JavaStyleIdentifier identifier = JavaStyleIdentifier.findFollowingIdentifier(expression, index);
                if(identifier == null){
                    throw new IllegalArgumentException("illegal expression, there is no entity following @.");
                }
                stringBuilder.append(expression, start, identifier.getStart()-DmlSetting.ENTITY_IDENTIFIER.length());
                if(!entityContainer.containsKey(identifier.getName())){
                    throw new IllegalArgumentException("entity not found: " + identifier.getName());
                }
                Object entityValue = entityContainer.get(identifier.getName());
                if(entityValue != null && entityValue instanceof String){
                    stringBuilder.append("'").append(entityValue).append("'");
                }else {
                    stringBuilder.append(entityValue);
                }
                start = identifier.getEnd();
            }
            stringBuilder.append(expression,start , expression.length());
            return stringBuilder.toString();
        }

        private Object evaluator(JobContext context, int parameterIndex) {
            Object result;
            if (dmlSetting.getParameter(parameterIndex) != null) {
                String expression = preEvaluate(context, dmlSetting.getParameter(parameterIndex));
                result = context.getExpressionEvaluator().evaluate(expression, context);
            } else {
                List<ColumnSetting> columns = DatabaseLoader.this.getSetting().getColumns();
                ColumnSetting column = columns.get(parameterIndex);
                result =
                        StringUtils.isEmpty(column.getExpression()) ? null :
                                context.getExpressionEvaluator().evaluate(column.getExpression(), context);
            }
            return result;
        }

        private void execute_preparedSql(JobContext context, Connection connection) throws SQLException {
            List<Object> parameters = new ArrayList<>();
            for (int i = 0; i < preparedStatement.getParameterMetaData().getParameterCount(); i++) {
                Object result = evaluator(context, i);
                DatabaseUtil.setObject(preparedStatement, i + 1, result);
                parameters.add(result);
            }
            if(log.isDebugEnabled()){
                log.debug("{" + this.dmlSetting.getId() + "} " +parameters);
            }
            preparedStatement.addBatch();
            batched++;
            try {
                if (batched % getSetting().getBatch() == 0) {
                    preparedStatement.executeBatch();
                    preparedStatement.clearParameters();
                    connection.commit();
                    count(batched);
                    batched = 0;
                    if (getLoadedRecordCount() % getCounterInterval(context) == 0)
                        log.info("loaded record [" + getLoadedRecordCount() + "] by " + getName());
                }
            } catch (SQLException e) {
                if (!this.dmlSetting.isTolerate())
                    throw e;
            }
        }

        private void execute_generatedSql(JobContext context, Connection connection) throws SQLException {
            createGeneratedSql();
            StringBuilder stringBuilder = new StringBuilder(valueSegment);
            for (int i = 0, markFrom = 0; i < preparedStatement.getParameterMetaData().getParameterCount(); i++) {
                Object result = evaluator(context, i);
                String value = toText(result);
                int index = stringBuilder.indexOf(DmlSetting.QUESTION_MARK, markFrom);
                stringBuilder.replace(index, index + 1, value);
                markFrom += value.length();
                int commaIndex = stringBuilder.indexOf(DmlSetting.COMMA_MARK, markFrom);
                if(commaIndex > markFrom){
                    markFrom = commaIndex;
                }
            }
            stringBuilder.append(",");
            generatedSql = generatedSql.concat(stringBuilder.toString());
            batched++;

            try {
                if (batched % getSetting().getBatch() == 0) {
                    // remove the last character ,
                    generatedSql = generatedSql.substring(0, generatedSql.length() - 1);
                    statement = connection.createStatement();
                    if(log.isDebugEnabled()){
                        log.debug("{" + this.dmlSetting.getId() + "} " +generatedSql);
                    }
                    statement.execute(generatedSql);
                    connection.commit();
                    count(batched);
                    generatedSql = null;
                    batched = 0;
                    if (getLoadedRecordCount() % getCounterInterval(context) == 0)
                        log.info("loaded record [" + getLoadedRecordCount() + "] by " + getName());
                }
            } catch (SQLException e) {
                if (!this.dmlSetting.isTolerate())
                    throw e;
            }
        }

        public void executeTail(Connection connection) throws SQLException {
            try {
                if (!this.dmlSetting.isMultiValue()) {
                    executeTail_preparedSql(connection);
                } else {
                    executeTail_generatedSql(connection);
                }
            } finally {
                close(this.statement);
                close(this.preparedStatement);
            }
        }


        private void executeTail_preparedSql(Connection connection) throws SQLException {
            try {
                if (this.preparedStatement != null && batched != 0) {
                    preparedStatement.executeBatch();
                    count(batched);
                    log.info("loaded record [" + getLoadedRecordCount() + "] by " + getName());
                    connection.commit();
                }
            } catch (SQLException e) {
                if (!this.dmlSetting.isTolerate())
                    throw e;
            }
        }

        private void executeTail_generatedSql(Connection connection) throws SQLException {
            if (StringUtils.isNotEmpty(generatedSql)) {
                // remove the last character ,
                generatedSql = generatedSql.substring(0, generatedSql.length() - 1);
                statement = connection.createStatement();
                statement.execute(generatedSql);
                connection.commit();
                count(batched);
                batched = 0;
                log.info("loaded record [" + getLoadedRecordCount() + "] by " + getName());
            }
        }

    }

    private static String toText(Object result) {
        String delim = "'";
        if(result == null)
            return "null";
        else if(result instanceof Timestamp){
            return delim + DateFormatUtils.format((Timestamp) result, "yyyy-MM-dd HH:mm:ss.SSS") + delim;
        }else if(result instanceof Date){
            return delim + DateFormatUtils.format((Date) result, "yyyy-MM-dd") + delim;
        }else if(result instanceof String){
            String text = StringUtils.replace((String) result, "'", "''");
            text = StringUtils.replace(text, "\\", "\\\\");
            return delim + text + delim;
        }
            return delim + result + delim;
    }

}
